import os
import sys
import inspect
import copy
import importlib
import re
import logging

import torch
import argparse

from config import *


def is_local_defined(object, member):
    srcfile = inspect.getsourcefile(object)
    return inspect.getsourcefile(member) == srcfile


def is_valid_Module(object, member):
    return inspect.isclass(member) and issubclass(member, torch.nn.Module) and is_local_defined(object, member)


def is_valid_function(object, member):
    return inspect.isfunction(member) and is_local_defined(object, member)


# find left (
def find_left_parentheses(string, right):
    assert string[right] == ')'
    stack = []
    for i in range(right, -1, -1):
        if string[i] == ')':
            stack.append(')')
        elif string[i] == '(':
            stack.pop()
            if len(stack) == 0:
                return i
    raise ValueError("{} should contain ()".format(string))


# find right )
def find_right_parentheses(string, left):
    stack = []
    for i in range(left, len(string)):
        if string[i] == '(':
            stack.append('(')
        elif string[i] == ')':
            stack.pop()
            if len(stack) == 0:
                return i
    raise ValueError("{} should contain ()".format(string))


def get_call_name(code, end):
    stack = []
    for i in range(end - 1, -1, -1):
        if code[i] in ["(", "[", "{"]:
            if stack:
                stack.pop()
            else:
                return code[i + 1:end], i + 1
        elif code[i] in [")", "]", "}"]:
            stack.append(code[i])
        elif stack:
            continue
        elif not (code[i].isalpha() or code[i].isdigit() or code[i] == '_' or code[i] == '.'):
            return code[i + 1:end], i + 1
    return "", -1


def convert_API(code, start, api_name=""):
    # 处理.shape(这样特殊的结构
    if api_name.startswith('.'):
        call_name, new_start = get_call_name(code, start)
        if start == -1 or call_name == "self":
            return code, start + 1
    else:
        call_name = api_name
        new_start = start

    # 找到完整的待转换的API. 例如: expr="nn.Conv2d(1,2,3)" args_str="(1,2,3)"
    left = code.find("(", start)
    assert left != -1
    right = find_right_parentheses(code, left)
    end = right
    expr = code[start:end + 1]
    args_str = code[left:right + 1]

    map_helper = ALL_MAPPING[api_name]
    new_expr = map_helper.convert(call_name, args_str)
    next_newline = code.find("\n", end + 1)
    fill_num = (expr.count("\n") - new_expr.count("\n"))
    if next_newline != -1:
        code = code[:new_start] + new_expr + code[end + 1:next_newline] + ("\n" * fill_num) + code[next_newline:]
    else:
        code = code[:new_start] + new_expr + ")" + ("\n" * fill_num) + code[end + 2:]
    return code, start + len(map_helper.ms_api.name)


def find_API(code, i):
    if code[i:].startswith("nn.") \
            or code[i:].startswith("F.") \
            or code[i:].startswith("torch.") \
            or code[i:].startswith('.'):
        j = code.find('(', i)
        if j != -1 and code[i:j] in ALL_TORCH_APIS:
            api_name = code[i:j]
            return api_name
    return ""


def convert_Function(fun_name, Fun, is_forward):
    _, line_no = inspect.getsourcelines(Fun)
    logging.info("line {:3d}: start converting function {}()".format(line_no, fun_name))

    code = inspect.getsource(Fun)
    code_saved = copy.copy(code)

    i = 0
    while i < len(code):
        api_name = find_API(code, i)
        if api_name:
            line_no1 = line_no + code[:i].count('\n')
            if api_name.startswith('.'):
                call_name, _ = get_call_name(code, i)
                if call_name in ["torch", "self"]:
                    i += 1
                    continue
            if api_name in ALL_MAPPING:
                if is_forward and api_name in NN_LIST:
                    warn_info = "line {:3d}: \t\tmindspore不支持在construct()中使用nn.xxx(): {}() -> {}()"
                    logging.warning(warn_info.format(line_no1, api_name, ALL_MAPPING[api_name].ms_api.name))
                    i += 1
                elif not is_forward and api_name in ALL_2P_LIST:
                    warn_info = "line {:3d}: \t\tmindspore不支持在非construct()代码段中使用原语: {} -> {}"
                    logging.warning(warn_info.format(line_no1, api_name, ALL_MAPPING[api_name].ms_api.name))
                    i += 1
                else:
                    code, i = convert_API(code, i, api_name)
                    logging.info("line {:3d}: \t\tconvert API success: {}".format(line_no1, api_name))
                continue
            else:
                warn_info = ". " + UNSUPPORTED_WARN_INFOS[api_name] if api_name in UNSUPPORTED_WARN_INFOS else ""
                logging.warning("line {:3d}: \t\tfound unsupported API: {}{}".format(line_no1, api_name, warn_info))
        i += 1
    return {code_saved: code} if code_saved != code else {}


def judge_forward(name, forward_list):
    is_forward = name in forward_list or name.split(".")[-1] == "forward"
    if is_forward:
        logging.debug("{} is a forward function".format(name))
    return is_forward


def convert_Module(module_name, Module, forward_list):
    _, line_no = inspect.getsourcelines(Module)
    logging.info("line {:3d}: start converting nn.Module {}".format(line_no, module_name))

    mapped = {}
    for name, member in inspect.getmembers(Module):
        if is_valid_function(Module, member):
            is_forward = judge_forward("{}.{}".format(module_name, name), forward_list)
            mapped.update(convert_Function(name, member, is_forward))
    return mapped


def get_forward_list(import_mod):
    forward_list = set()

    ext_funs = []
    for name, member in inspect.getmembers(import_mod):
        if is_valid_function(import_mod, member):
            ext_funs.append(name)

    for name, member in inspect.getmembers(import_mod):
        if is_valid_Module(import_mod, member):
            Module = member
            self_funs = []
            for name1, member1 in inspect.getmembers(Module):
                if name1 == "forward":
                    forward_fun = member1
                    forward_list.add(name + '.' + name1)
                elif is_valid_function(Module, member1):
                    self_funs.append(name1)

            forward_code = inspect.getsource(forward_fun)

            for fun in self_funs:
                if "self." + fun in forward_code:
                    forward_list.add(name + '.' + fun)

            for fun in ext_funs:
                if fun in forward_code:
                    forward_list.add(fun)

    return forward_list


def convert(in_file, out_file):
    assert os.path.isfile(in_file)
    assert in_file.endswith(".py")
    full_name = os.path.abspath(in_file)
    print("Start converting {}".format(full_name))
    import_path = os.path.abspath(os.path.join(os.path.dirname(in_file), ".."))
    import_name = os.path.dirname(full_name).split('/')[-1] + '.' + full_name.split('/')[-1][:-len(".py")]
    logging.debug("add '{}' to sys.path".format(import_path))
    sys.path.append(import_path)
    logging.debug("start import '{}'".format(import_name))
    import_mod = importlib.import_module(import_name)

    # 得到forward函数及其调用函数的的列表
    forward_list = get_forward_list(import_mod)
    logging.debug("forward_list: {}".format(forward_list))

    # 对nn.Modlue/python函数进行替换
    mapping = {}
    tasks = []
    for name, member in inspect.getmembers(import_mod):
        if is_valid_Module(import_mod, member):
            _, line_no = inspect.getsourcelines(member)
            tasks.append((line_no, convert_Module, (name, member, forward_list)))
        elif is_valid_function(import_mod, member):
            _, line_no = inspect.getsourcelines(member)
            is_forward = judge_forward("{}.{}".format(import_mod, name), forward_list)
            tasks.append((line_no, convert_Function, (name, member, is_forward)))
    tasks.sort()
    for _, convert_fun, args in tasks:
        mapping.update(convert_fun(*args))
    code = inspect.getsource(import_mod)
    for k, v in mapping.items():
        code = code.replace(k, v)

    # 全局替换
    code = 'import mindspore\n' \
           'import mindspore.nn as nn\n' \
           'import mindspore.ops.operations as P\n' \
           + code
    code = re.sub(r"import torch", "# import torch", code)
    code = re.sub(r"from torch", "# from torch", code)
    code = re.sub(r"\(nn.Module\):", "(nn.Cell):", code)
    code = re.sub(r"def\sforward\(self", "def construct(self", code)
    code = re.sub(r"nn.Linear", "nn.Dense", code)
    code = re.sub(r"\(nn.Sequential\)", "(nn.SequentialCell)", code)
    code = re.sub(r"nn.Linear", "nn.Dense", code)
    code = re.sub(r'nn.init.', 'pass  # nn.init.', code)

    with open(out_file, 'w') as f:
        f.write(code)
    print("Convert success! result is written to {}".format(out_file))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="A script converter from PyTorch to MindSpore")
    parser.add_argument("--in_file", required=True, help="the name of the file to convert")
    parser.add_argument("--out_file", default="out.py", help="the output filename (default: out.py).")
    parser.add_argument("--print_all", help="print full log to stdout (default: False).", action="store_true")
    args = parser.parse_args()

    logger = logging.getLogger()
    logging.basicConfig(level=logging.DEBUG, format='%(levelname)-8s%(message)s')
    if args.print_all:
        logger.setLevel(logging.INFO)
    else:
        logger.setLevel(logging.WARNING)

    convert(args.in_file, args.out_file)
